import json
import logging
import re
from collections import deque

from sonetto.agents.base import BaseAgent
from sonetto.agents.chat import ChatAgent
from sonetto.core.message import Message
from sonetto.core.scheduler import Scheduler

base_url_glm = "https://open.bigmodel.cn/api/paas/v4"
api_key_glm = "ff37fbba1e80d4cb5f38ac819893f4a1.c6T0Zm3Sd9YODRLV"
model_glm = "glm-4-flash"

base_url_ds = "https://api.deepseek.com"
api_key_ds = "sk-3bb4d61950df4f4d833f800848e1f461"
model_ds = "deepseek-chat"

LARGE_BOUNDARY = 700
LINE_PER_CHUNK = 50

class SplitNode(BaseAgent):
    """文本分块节点"""

    def __init__(self, agent_id: str, duty: str, line_per_chunk: int = 70):
        super().__init__(agent_id=agent_id, duty=duty)
        self.queue = deque()
        self.line_per_chunk = line_per_chunk

    def init(self, file_path: str):
        # './docs/【03-规章制度】/龙腾集团员工休假管理办法（2022版）签批版.md'
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.readlines()

            # 每 line_per_chunk 行分割一次
            for i in range(0, len(content), self.line_per_chunk):
                chunk = ''.join(content[i:i + self.line_per_chunk])
                self.queue.append(chunk)

    def forward(self, message: Message) -> Message:
        if message.sender == 'user':
            self.init(message.content)
            message.metadata['file_path'] = message.content

        if self.queue:
            chunk = self.queue.popleft()
            return Message(
                content=
                f'''接下来请将下面的文本进行清洗。去除冗余内容的同时，保持有效内容的原文表述，不要擅自对有效内容进行简化。必要时补充上下文。
```markdown
{chunk}
```
''',
                sender=self.agent_id,
                metadata=message.metadata)
        else:
            return Message(content='''分块完成！！！''', sender=self.agent_id)


class KnowledgeBaseNode(BaseAgent):
    """知识入库节点"""

    def __init__(self, agent_id: str, duty):
        super().__init__(agent_id=agent_id, duty=duty)

    def forward(self, message: Message) -> Message:
        json_str = re.search(r'```json\n(.*?)\n```', message.content,
                             re.DOTALL)
        if json_str and message.sender == 'checker':
            json_str = json_str.group(1)
            large_chunk = []

            # 将 json 字符串转换为字典列表
            chunks = json.loads(json_str)
            for chunk in chunks:
                with open(message.metadata['file_path'].replace('.md', '.txt'),
                          'a',
                          encoding='utf-8') as f:
                    f.write(str(chunk) + '\n')
                if len(chunk['content']) > LARGE_BOUNDARY:
                    large_chunk.append(chunk['content'])

            if large_chunk:
                large_chunk_txt = "\n\n".join(large_chunk)
                return Message(content=f'''请对以下文段进行命题分割：

```markdown
{large_chunk_txt} 
```
''',
                               sender=self.agent_id,
                               receiver='propositioner',
                               metadata=message.metadata)
            else:
                return Message(content='',
                               sender=self.agent_id,
                               receiver='spliter',
                               metadata=message.metadata)
        elif json_str and message.sender == 'propositioner':
            propositions = json.loads(json_str.group(1))
            for proposition in propositions:
                with open(message.metadata['file_path'].replace('.md', '.txt'),
                          'a',
                          encoding='utf-8') as f:
                    f.write(proposition + '\n')
            return Message(content='',
                           sender=self.agent_id,
                           receiver='spliter',
                           metadata=message.metadata)


washer = ChatAgent(agent_id="washer",
                   base_url=base_url_ds,
                   api_key=api_key_ds,
                   model=model_ds,
                   system_prompt='''### 文本清洗指南：
1. 清洗文本中的广告、目录、口号、页码。
2. 清洗文本中的无用字符，例如多余空格、换行符、引号。
3. 保留markdown格式的图片、将表格信息提取为自然语言分点列项。
4. 将文本样式调整为易于大模型阅读并进行文本分块的内容。
5. 除了上述内容，**不要擅自修改文本内容**，保持原样输出。
6. 如果用户给你的整段文本中没有有效的内容，请告诉 chunker。
7. 用户的文本可能并不完整，可能会有截断的情况，属于正常情况。
如果上一轮交给 chunker 的内容不完整，可以在当前轮次适当**结合上下文对文本内容进行补充**。

文本清洗完成后，请将清洗后的文本交给 chunker 进行文本分块。
''',
                   duty='清洗文本，将清洗后的文本交接给 chunker 进行文本分块。',
                   max_num_context=5,
                   use_protocol=True)

chunker = ChatAgent(agent_id="chunker",
                    base_url=base_url_ds,
                    api_key=api_key_ds,
                    model=model_ds,
                    system_prompt='''你划分的块最后会作为问答系统的参考内容，故其语义的独立完整性非常重要。

### 文本分块划分指南：
1. 各文本块应在主题/语义上连贯且独立，确保它们脱离上下文也能被理解。
2. 语义边界划分：
    - 分节边界（章节、标题）
    - 内容主题转换
    - 形成逻辑单元的段落组
    - 如果看到问答对，请直接将问答对划分为一个块。
3. **图片或表格不能单独划分为一块**，应当和最近的文本块合并。
4. 如果文本块的内容过长，超过 1400 字符，请按语义将其划分为多个块。

分块过程中，你必须保持文本的原样输出，**不要擅自修改文本内容**。

如果用户给的分块文本因为截断导致最后的内容无法组成完整的分块，**请不要将其划分为一个块**，而是等待下一个分块的内容到来后再进行划分。
```
''',
                    duty='识别文本的语义边界，将文本划分为多个块。将划分后的块交接给 checker 进行最后的校验。',
                    max_num_context=6,
                    use_protocol=True)

checker = ChatAgent(agent_id="checker",
                    base_url=base_url_ds,
                    api_key=api_key_ds,
                    model=model_ds,
                    system_prompt='''### 文本校验指南：
1. 校验文本块的语义完整性和独立性。下面是文本分块指南：
    - 各文本块应在主题/语义上连贯且独立。
    - 语义边界划分：
        - 分节边界（章节、标题）
        - 内容主题转换
        - 形成逻辑单元的段落组
        - 问答对直接划分为一个块。
        - 图片和表格划到哪一块都行，不影响分块校验结果。
2. 如果文本块的语义不完整或不独立，请指导 chunker 重新调整分块结果。
3. 如果校验通过，提取 chunker 的分块结果为 jsonl 对象，即每个分块提取为一个单行 jsonl 对象，包含 index 和 content 两个字段。index 是分块摘要，content 是分块内容。
示例如下:
```json
[
{"index": "员工权益-年假", "content": "员工有权享受年假。\\n年假是指员工在工作满一年后，享有的带薪休假。"},
{"index": "龙腾集团规章制度-保密制度", "content": "保密的相关内容..."}
]
```
请注意转义字符的使用，例如换行符、引号等。你的分块结果将会被 JSON 解析器解析。
确保你的分块结果是有效的 JSON 格式，并保证你的分块结果包裹在 ```json 代码块中。

若校验通过，请直接将分块结果交给 manager，无需给 chunker 反馈。

''',
                    duty='校验文本块的语义完整性和独立性，并负责提取分块结果交给 manager。',
                    max_num_context=5,
                    use_protocol=True)

propositioner = ChatAgent(agent_id="propositioner",
                          base_url=base_url_ds,
                          api_key=api_key_ds,
                          model=model_ds,
                          system_prompt="""
你将收到很多文本段落。将文本内容分解为清晰简单的命题，确保它们脱离上下文也能被理解。这些命题将用于问答系统的知识库中，因此它们应该是独立、完整的，并且不依赖于上下文。

1. 将复合句拆分为简单句，但结果中不应出现前缀相同，后缀不同的句子。尽可能保留输入中的原始表述，但也确保每个命题至少包含主谓宾。

2. 对于任何伴随有额外描述信息的命名实体，将这些信息分离成独立的命题。

3. 通过给名词或整个句子添加必要的修饰语，以及将代词（如"它"、"他"、"她"、"他们"、"这个"、"那个"）替换为它们所指代实体的全名，使命题去语境化。

4. 去除缺少有效信息的命题，如目录、广告词以及其他不可能作为某个问题答案的命题。

5. 将所有文本段落全部拆为命题，并将结果组合为JSON格式的字符串列表呈现。

输入: 
```markdown
关于复活节兔(Osterhase)的最早证据是由医学教授乔治·弗兰克·冯·弗兰肯瑙于1678年在德国西南部记录的，但直到18世纪它在德国其他地区仍不为人所知。

学者理查德·瑟蒙写道："兔子经常在春天的花园中被发现，因此可能成为解释隐藏在那里供儿童寻找的彩蛋起源的一个便利解释。或者，有一种欧洲传统认为兔子会下蛋，因为兔子的窝和田凫鸟的巢看起来非常相似，两者都出现在草地上，并且都是在春天首次被发现。在十九世纪，复活节贺卡、玩具和书籍的影响使得复活节兔在整个欧洲变得流行。德国移民随后将这一习俗传播到英国和美国，在那里它演变成了复活节兔。"
```

输出: 
```json
[ "关于复活节兔的最早证据是由乔治·弗兰克·冯·弗兰肯瑙于1678年在德国西南部记录的。", "乔治·弗兰克·冯·弗兰肯瑙是一位医学教授。", "复活节兔的证据直到18世纪在德国其他地区仍不为人所知。", "理查德·瑟蒙是一位学者。", "理查德·瑟蒙写出了一个关于兔子与复活节传统之间联系的可能解释的假设。", "兔子经常在春天的花园中被发现。", "兔子可能成为解释隐藏在花园中供儿童寻找的彩蛋起源的一个便利解释。", "有一种欧洲传统认为兔子会下蛋。", "兔子的窝和田凫鸟的巢看起来非常相似。", "兔子和田凫鸟的巢都出现在草地上，并且都是在春天首次被发现。", "在十九世纪，复活节贺卡、玩具和书籍的影响使得复活节兔在整个欧洲变得流行。", "德国移民将复活节兔的习俗传播到英国和美国。", "复活节兔的习俗在英国和美国演变成了复活节兔。" ]
```
""",
                          max_num_context=3)

scheduler = Scheduler()
spliter = SplitNode(agent_id="spliter",
                    duty="文档分割节点，负责将长文本分割为多个块。",
                    line_per_chunk=LINE_PER_CHUNK)
manager = KnowledgeBaseNode(agent_id="manager", duty="管理员，负责完成最后的知识入库工作。")
scheduler.register_agent(spliter)
scheduler.register_agent(washer)
scheduler.register_agent(chunker)
scheduler.register_agent(checker)
scheduler.register_agent(manager)
scheduler.register_agent(propositioner)
scheduler.set_route(spliter.agent_id, washer.agent_id)
scheduler.set_route(propositioner.agent_id, manager.agent_id)

response = scheduler.run_stream(
    task='./docs/【03-规章制度】/龙腾出行职业发展通道管理办法 2024.md',
    cancel_token='分块完成',
    cancel_from=[spliter.agent_id],
    max_step=414,
)
